import tkinter as tk
from PIL import ImageGrab
import cv2
import numpy as np
import os
import threading
import time
import pyautogui

class ScreenRecorder:
    def __init__(self):
        self.root = tk.Tk()
        self.root.title("录屏工具")
        self.root.geometry("400x300")
        
        # 当前目录作为存储路径
        self.save_dir = os.getcwd()
        self.recording = False
        self.start_x = None
        self.start_y = None
        self.rect = None
        self.canvas = None
        self.video_writer = None
        
        # GUI组件
        self.label = tk.Label(self.root, text="录屏工具 v1.0")
        self.label.pack(pady=10)
        
        self.select_btn = tk.Button(self.root, text="选择录制区域", command=self.start_selection)
        self.select_btn.pack(pady=10)
        
        self.start_btn = tk.Button(self.root, text="开始录制", command=self.start_recording, state="disabled")
        self.start_btn.pack(pady=10)
        
        self.stop_btn = tk.Button(self.root, text="停止录制", command=self.stop_recording, state="disabled")
        self.stop_btn.pack(pady=10)
        
        self.status = tk.Label(self.root, text="请先选择录制区域")
        self.status.pack(pady=10)
        
        self.root.mainloop()

    def get_unique_filename(self):
        """生成唯一的文件名，避免覆盖"""
        base_name = "screen_record"
        extension = ".mp4"
        counter = 0
        filename = os.path.join(self.save_dir, f"{base_name}{extension}")
        
        while os.path.exists(filename):
            counter += 1
            filename = os.path.join(self.save_dir, f"{base_name}_{counter}{extension}")
        
        return filename

    def start_selection(self):
        """开始选择录制区域"""
        self.root.withdraw()
        
        self.selection_window = tk.Toplevel()
        self.selection_window.attributes('-fullscreen', True)
        self.selection_window.attributes('-alpha', 0.3)
        
        self.canvas = tk.Canvas(self.selection_window, cursor="cross")
        self.canvas.pack(fill="both", expand=True)
        
        self.canvas.bind("<ButtonPress-1>", self.on_press)
        self.canvas.bind("<B1-Motion>", self.on_drag)
        self.canvas.bind("<ButtonRelease-1>", self.on_release)
        self.selection_window.bind("<Escape>", lambda e: self.cancel_selection())

    def on_press(self, event):
        """鼠标按下，记录起点"""
        self.start_x = self.canvas.winfo_pointerx()
        self.start_y = self.canvas.winfo_pointery()
        self.rect = self.canvas.create_rectangle(self.start_x, self.start_y, self.start_x, self.start_y, outline="red", width=2)

    def on_drag(self, event):
        """鼠标拖动，更新矩形区域"""
        cur_x = self.canvas.winfo_pointerx()
        cur_y = self.canvas.winfo_pointery()
        self.canvas.coords(self.rect, self.start_x, self.start_y, cur_x, cur_y)

    def on_release(self, event):
        """鼠标释放，确定录制区域"""
        self.end_x = self.canvas.winfo_pointerx()
        self.end_y = self.canvas.winfo_pointery()
        
        self.left = min(self.start_x, self.end_x)
        self.top = min(self.start_y, self.end_y)
        self.right = max(self.start_x, self.end_x)
        self.bottom = max(self.start_y, self.end_y)
        
        self.selection_window.destroy()
        self.root.deiconify()
        self.status.config(text="区域已选择，点击‘开始录制’")
        self.start_btn.config(state="normal")

    def cancel_selection(self):
        """取消选择"""
        self.selection_window.destroy()
        self.root.deiconify()
        self.status.config(text="请先选择录制区域")
        self.start_btn.config(state="disabled")

    def start_recording(self):
        """开始录制"""
        self.recording = True
        self.start_btn.config(state="disabled")
        self.stop_btn.config(state="normal")
        self.status.config(text="正在录制...")
        
        # 获取屏幕分辨率和帧率
        self.width = self.right - self.left
        self.height = self.bottom - self.top
        self.fps = 20  # 可调整帧率
        
        # 初始化视频写入器
        fourcc = cv2.VideoWriter_fourcc(*"mp4v")
        self.save_path = self.get_unique_filename()
        self.video_writer = cv2.VideoWriter(self.save_path, fourcc, self.fps, (self.width, self.height))
        
        # 异步录制
        self.record_thread = threading.Thread(target=self.record_screen)
        self.record_thread.start()

    def record_screen(self):
        """录制屏幕"""
        while self.recording:
            # 捕获屏幕区域
            img = ImageGrab.grab(bbox=(self.left, self.top, self.right, self.bottom))
            frame = np.array(img)
            frame = cv2.cvtColor(frame, cv2.COLOR_RGB2BGR)  # 转换为OpenCV格式
            
            # 写入视频
            self.video_writer.write(frame)
            time.sleep(1 / self.fps)  # 控制帧率
        
        # 结束录制
        self.video_writer.release()

    def stop_recording(self):
        """停止录制"""
        self.recording = False
        self.stop_btn.config(state="disabled")
        self.start_btn.config(state="normal")
        self.status.config(text=f"录制完成，保存为 {os.path.basename(self.save_path)}")

if __name__ == "__main__":
    ScreenRecorder()